package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"sort"
	"strconv"
	"strings"
)

// Stats 结构体用于存储每个前缀的统计信息，包括TTL=-1的计数
type Stats struct {
	Count       int64 // 总数量
	Size        int64 // 总大小
	TTLMinusOne int64 // TTL为-1的数量
}

// StatMap 是一个映射，key为前缀，value为对应的统计数据
type StatMap map[string]*Stats

func main() {
	// 定义命令行参数
	csvFile := flag.String("file", "", "Path to the CSV file generated by rdb-tools (e.g., memory.csv)")
	keyPrefixDelimiter := flag.String("delimiter", "_:", "Delimiter(s) to split the Redis key for hierarchy (default: '_', supports multiple like '_:')")
	sizeColumnName := flag.String("size-column", "size_in_bytes", "Name of the column containing the size in the CSV (default: 'size_in_bytes')")
	keyColumnName := flag.String("key-column", "key", "Name of the column containing the key in the CSV (default: 'key')")
	targetLevel := flag.Int("level", 3, "Specific hierarchy level to analyze (default: 3)")
	topN := flag.Int("top", 50, "Number of top results to display (default: 50, use -1 for all)")

	flag.Parse()

	if *csvFile == "" {
		log.Fatal("Please provide the path to the CSV file using the -file flag.")
	}

	if *targetLevel < 1 {
		log.Fatal("Target level (-level) must be a positive integer.")
	}

	if *topN < -1 {
		log.Fatal("Top N (-top) must be -1 (for all) or a positive integer.")
	}

	// 打开CSV文件
	file, err := os.Open(*csvFile)
	if err != nil {
		log.Fatalf("Failed to open CSV file: %v", err)
	}
	defer func(file *os.File) {
		err := file.Close()
		if err != nil {
			log.Printf("Failed to close file: %v", err)
		}
	}(file)

	// 创建CSV Reader
	reader := csv.NewReader(file)

	// 读取CSV头部
	header, err := reader.Read()
	if err != nil {
		log.Fatalf("Failed to read CSV header: %v", err)
	}

	// 在头部中查找key, size, 和 expiry 列的索引
	keyIndex := -1
	sizeIndex := -1
	expiryIndex := -1
	for i, columnName := range header {
		// 处理可能的BOM (Byte Order Mark)
		columnName = strings.TrimPrefix(columnName, "\uFEFF")
		if columnName == *keyColumnName {
			keyIndex = i
		} else if columnName == *sizeColumnName {
			sizeIndex = i
		} else if columnName == "expiry" {
			expiryIndex = i
		}
	}

	if keyIndex == -1 {
		log.Fatalf("Could not find column '%s' in the CSV header.", *keyColumnName)
	}
	if sizeIndex == -1 {
		log.Fatalf("Could not find column '%s' in the CSV header.", *sizeColumnName)
	}

	// 初始化存储指定层级统计数据的映射
	targetLevelStats := make(StatMap)

	// 逐行读取CSV数据
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Printf("Error reading CSV record: %v", err)
			continue // 跳过错误行
		}

		// 获取key和size
		key := record[keyIndex]
		sizeStr := record[sizeIndex]

		// 转换size为整数
		size, err := strconv.ParseInt(sizeStr, 10, 64)
		if err != nil {
			log.Printf("Error parsing size for key '%s': %v. Skipping record.", key, err)
			continue
		}

		// 初始化TTL相关变量
		isTTLMinusOne := expiryIndex != -1 && (record[expiryIndex] == "" || record[expiryIndex] == "-1")

		// 按照指定分隔符分割key
		parts := splitByDelimiters(key, *keyPrefixDelimiter)

		// 检查key是否有足够的部分来达到目标层级
		if len(parts) < *targetLevel {
			// 如果key的层级少于目标层级，则跳过
			continue
		}

		// 构建目标层级的前缀
		// 使用原始分隔符参数中的第一个分隔符来连接前缀
		delimiter := getFirstDelimiter(*keyPrefixDelimiter)
		prefix := strings.Join(parts[:*targetLevel], delimiter)

		// 获取或创建当前前缀的统计数据
		if targetLevelStats[prefix] == nil {
			targetLevelStats[prefix] = &Stats{}
		}

		// 更新统计数据
		targetLevelStats[prefix].Count++
		targetLevelStats[prefix].Size += size
		if isTTLMinusOne {
			targetLevelStats[prefix].TTLMinusOne++
		}
	}

	// 输出指定层级的结果，包含TTL=-1统计
	fmt.Printf("Analysis Results for Level %d (Including TTL=-1 Count):\n", *targetLevel)
	fmt.Println("=====================================================")
	if len(targetLevelStats) == 0 {
		fmt.Println("No data found for the specified level.")
		return
	}

	// 打印表头
	fmt.Printf("%-50s | %10s | %15s | %10s\n", "Prefix", "Total", "Total Size (MB)", "TTL=-1")
	fmt.Println(strings.Repeat("-", 100))

	// 为了输出有序，对 keys 进行排序
	keys := make([]string, 0, len(targetLevelStats))
	for k := range targetLevelStats {
		keys = append(keys, k)
	}

	// 按数量排序（从高到低）
	sort.Slice(keys, func(i, j int) bool {
		return targetLevelStats[keys[i]].Count > targetLevelStats[keys[j]].Count
	})

	// Simple string sort (uncomment import "sort" and sort.Strings(keys) for alphabetical)
	// sort.Strings(keys)

	// 打印统计数据
	rowsToPrint := *topN
	if *topN == -1 {
		// 显示所有行
		rowsToPrint = len(keys)
	} else if *topN > len(keys) {
		// 如果要求显示的行数超过实际行数，则显示所有行
		rowsToPrint = len(keys)
	}

	// 实际打印指定的行数
	for i, prefix := range keys {
		if i >= rowsToPrint {
			break
		}
		stat := targetLevelStats[prefix]
		fmt.Printf("%-50s | %10d | %15.2f | %10d\n", prefix, stat.Count, float64(stat.Size)/1024/1024, stat.TTLMinusOne)
	}

	// 如果不是显示所有行，提示用户
	if *topN != -1 && *topN < len(keys) {
		fmt.Printf("\n... %d more entries not shown (use -top=%d or -top=-1 to show all) ...\n", len(keys)-*topN, len(keys))
	}
	// 如果不关心顺序，可以直接遍历 map
	// for prefix, stat := range targetLevelStats {
	//     fmt.Printf("%-50s | %10d | %15d | %15d\n", prefix, stat.Count, stat.Size, stat.TTLMinusOne)
	// }
}

// splitByDelimiters 使用指定的分隔符分割字符串，支持多个分隔符
func splitByDelimiters(s string, delimiters string) []string {
	// 如果只有一个分隔符，直接使用strings.Split
	if len(delimiters) == 1 {
		return strings.Split(s, delimiters)
	}

	// 创建一个用于替换的映射
	replacer := strings.NewReplacer()

	// 使用第一个分隔符作为统一的分隔符
	primaryDelimiter := string(delimiters[0])

	// 为所有其他分隔符创建替换规则
	for _, r := range delimiters {
		if string(r) != primaryDelimiter {
			replacer = strings.NewReplacer(string(r), primaryDelimiter)
		}
	}

	// 标准化字符串
	normalized := replacer.Replace(s)

	// 分割并过滤空字符串
	parts := strings.Split(normalized, primaryDelimiter)

	// 移除空字符串
	var result []string
	for _, part := range parts {
		if part != "" {
			result = append(result, part)
		}
	}

	return result
}

// getFirstDelimiter 获取分隔符字符串中的第一个分隔符
func getFirstDelimiter(delimiters string) string {
	if len(delimiters) == 0 {
		return "_"
	}
	return string(delimiters[0])
}
